package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/go-redis/redis/v8"
	"github.com/streadway/amqp"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
)

const timeout = 5 * time.Second // 超时时间设置为 2 秒

// 打印用法和帮助信息
func printUsageAndExit() {
	fmt.Println("用法:")
	fmt.Println("./conn_test <服务类型> -h <主机地址> [-P <端口>] [-u <用户名>] [-p <密码>]")

	fmt.Println("服务类型:")
	fmt.Println("  mysql  : 测试 MySQL 连接")
	fmt.Println("  redis  : 测试 Redis 连接")
	fmt.Println("  es     : 测试 Elasticsearch 连接")
	fmt.Println("  mq     : 测试 RabbitMQ 连接")
	fmt.Println("  mongo  : 测试 MongoDB 连接")

	fmt.Println("示例:")
	fmt.Println("  ./conn_test mysql -h 127.0.0.1 -P 3306 -u <用户> -p <密码>")
	fmt.Println("  ./conn_test redis -h 127.0.0.1 -P 6379 -p <密码>")

	os.Exit(1)
}

// 检查端口是否通
func checkPort(host string, port int) bool {
	address := fmt.Sprintf("%s:%d", host, port)
	_, err := net.DialTimeout("tcp", address, timeout)
	if err != nil {
		fmt.Printf("端口 %d 无法连接: %s\n", port, err)
		return false
	}
	fmt.Printf("端口 %d 畅通!!!\n", port)
	return true
}

// MySQL 测试
func testMySQL(user, password, host string, port int) {
	if !checkPort(host, port) {
		return
	}

	// 配置 MySQL 数据源名称 (DSN)
	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/?timeout=2s", user, password, host, port)

	// 使用 GORM 打开数据库连接
	db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{})
	if err != nil {
		fmt.Printf("MySQL 连接失败: %s\n", err)
		return
	}
	fmt.Println("MySQL 连接成功!!!")

	// 连接成功，执行 SHOW GRANTS 命令
	var result string
	err = db.Raw("SHOW GRANTS;").Scan(&result).Error
	if err != nil {
		fmt.Printf("执行 SHOW GRANTS 失败: %s\n", err)
	} else {
		fmt.Println("用户权限：", result)
	}
}

func testRedis(user, password, host string, port int) {
	if !checkPort(host, port) {
		return
	}

	// 创建 Redis 客户端
	rdb := redis.NewClient(&redis.Options{
		Addr:     fmt.Sprintf("%s:%d", host, port),
		Username: user,
		Password: password, // 使用密码认证
		DB:       0,        // 使用默认数据库
	})

	ctx := context.Background()
	// 创建上下文并设置超时
	_, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	// 使用 Ping 来测试连接
	_, err := rdb.Ping(ctx).Result() // 使用 ctx 进行连接
	if err != nil {
		fmt.Printf("Redis 连接失败: %s\n", err)
		return
	}
	fmt.Println("Redis 连接成功!!!")

	// 执行 INFO replication 命令，传入 ctx 和 section 参数
	info, err := rdb.Info(ctx, "replication").Result()
	if err != nil {
		fmt.Printf("执行 INFO replication 失败: %s\n", err)
		return
	}

	// 打印 INFO replication 的第一行内容
	lines := strings.Split(info, "\n") // 按行分割信息
	if len(lines) > 0 {
		fmt.Printf("Redis 复制信息：%v\n", lines[1])

	} else {
		fmt.Println("INFO replication 没有返回有效数据")
	}
}

// Elasticsearch 测试
func testElasticsearch(user, password, host string, port int) {
	if !checkPort(host, port) {
		return
	}
	var httpURL string
	var httpsURL string

	if user != "" && password != "" {
		httpURL = fmt.Sprintf("http://%s:%s@%s:%d", user, password, host, port)
		httpsURL = fmt.Sprintf("https:///%s:%s@%s:%d", user, password, host, port)
	} else {
		httpURL = fmt.Sprintf("http://%s:%d", host, port)
		httpsURL = fmt.Sprintf("https://%s:%d", host, port)
	}

	var Status1 int
	var Status2 int
	resp, err := http.Get(httpURL)
	if err != nil {
		Status1 = 0
	} else {
		defer resp.Body.Close()
		if resp.StatusCode == http.StatusOK && containsClusterName(resp.Body) {
			Status1 = 1
		}
	}

	// 测试 HTTPS 协议
	resp, err = http.Get(httpsURL)
	if err != nil {
		Status2 = 0
	} else {
		defer resp.Body.Close()
		if resp.StatusCode == http.StatusOK && containsClusterName(resp.Body) {
			Status2 = 1
		}
	}
	if Status1 == 1 || Status2 == 1 {

		fmt.Println("Elasticsearch 连接成功")
	} else {
		fmt.Println("Elasticsearch 连接失败")
	}

}

// RabbitMQ 测试
func testRabbitMQ(user, password, host string, port int) {
	if !checkPort(host, port) {
		return
	}

	url := fmt.Sprintf("amqp://%s:%s@%s:%d/", user, password, host, port)

	_, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	conn, err := amqp.DialConfig(url, amqp.Config{
		Dial: func(network, addr string) (net.Conn, error) {
			return net.DialTimeout(network, addr, timeout)
		},
	})
	if err != nil {
		fmt.Printf("RabbitMQ 连接失败: %s\n", err)
	} else {
		defer conn.Close()
		fmt.Println("RabbitMQ 连接成功")
	}
}

func testMongoDB(user, password, host string, port int) {
	if !checkPort(host, port) {
		return
	}

	// 构建 MongoDB URI
	uri := fmt.Sprintf("mongodb://%s:%s@%s:%d/?connectTimeoutMS=2000", user, password, host, port)

	// 使用 mongo.Connect 而不是 mongo.NewClient
	client, err := mongo.Connect(context.Background(), options.Client().ApplyURI(uri))
	if err != nil {
		fmt.Printf("MongoDB 客户端创建失败: %s\n", err)
		return
	}
	defer client.Disconnect(context.Background()) // 确保断开连接

	// 创建上下文，并设置超时时间
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	// 尝试连接到 MongoDB 实例
	err = client.Ping(ctx, nil)
	if err != nil {
		fmt.Printf("MongoDB 连接失败: %s\n", err)
	} else {
		fmt.Println("MongoDB 连接成功")

		// 连接成功后，执行 show dbs 操作
		databases, err := client.ListDatabaseNames(ctx, bson.M{})
		if err != nil {
			fmt.Printf("获取数据库列表失败: %s\n", err)
		} else {
			fmt.Println("可用的数据库：")
			for _, db := range databases {
				fmt.Println(db)
			}
		}
	}
}

// 检查响应体是否包含 cluster_name
func containsClusterName(body io.Reader) bool {
	var jsonResponse map[string]interface{}
	decoder := json.NewDecoder(body)
	err := decoder.Decode(&jsonResponse)
	if err != nil {
		fmt.Printf("解析响应体失败: %s\n", err)
		return false
	}

	// 检查是否包含 cluster_name 字段
	if clusterName, exists := jsonResponse["cluster_name"]; exists && clusterName != nil {
		fmt.Printf("连接成功，集群名称: %s\n", clusterName)
		return true
	}
	return false
}
func main() {
	// 检查命令行参数
	if len(os.Args) < 3 {
		fmt.Println("错误: 参数不足")
		printUsageAndExit()
	}

	service := os.Args[1]
	flagSet := flag.NewFlagSet(service, flag.ExitOnError)

	host := flagSet.String("h", "", "主机地址")
	port := flagSet.Int("P", 0, "端口号")
	user := flagSet.String("u", "", "用户名")
	password := flagSet.String("p", "", "密码")

	flagSet.Parse(os.Args[2:])

	// 检查是否提供了主机地址
	if *host == "" {
		fmt.Println("错误: 主机地址 (-h) 必须提供")
		printUsageAndExit()
	}

	// 根据服务类型调用对应的测试函数
	switch service {
	case "mysql":
		if *port == 0 {
			*port = 3306
		}
		testMySQL(*user, *password, *host, *port)
	case "redis":
		if *port == 0 {
			*port = 6379
		}
		testRedis(*user, *password, *host, *port)
	case "es":
		if *port == 0 {
			*port = 9200
		}
		testElasticsearch(*user, *password, *host, *port)
	case "mq":
		if *port == 0 {
			*port = 5672
		}
		testRabbitMQ(*user, *password, *host, *port)
	case "mongo":
		if *port == 0 {
			*port = 27017
		}
		testMongoDB(*user, *password, *host, *port)
	default:
		fmt.Printf("错误: 未知的服务类型 '%s'\n", service)
		printUsageAndExit()
	}
}
